#!/usr/bin/env python3
"""
AES加密测试脚本 - 为 /users/send_sms_code/ 接口生成加密数据
"""
import json
import base64
import time
import sys

try:
    from Crypto.Cipher import AES
    from Crypto.Util.Padding import pad
    from Crypto.Random import get_random_bytes
except ImportError:
    print("❌ 缺少依赖包，请安装: pip install pycryptodome")
    sys.exit(1)

class EncryptionManager:
    def __init__(self, secret_key: str = "qwertyuiop1234567654zaqwsxcderfv"):
        self.secret_key = secret_key.encode('utf-8')[:32].ljust(32, b'\0')
    
    def aes_encrypt(self, data: str) -> str:
        """AES256加密"""
        iv = get_random_bytes(16)
        cipher = AES.new(self.secret_key, AES.MODE_CBC, iv)
        padded_data = pad(data.encode('utf-8'), AES.block_size)
        encrypted = cipher.encrypt(padded_data)
        return base64.b64encode(iv + encrypted).decode('utf-8')

def create_encrypted_payload(data: dict) -> dict:
    """创建加密载荷"""
    # 添加时间戳防重放攻击
    data["timestamp"] = int(time.time())
    
    # 转换为JSON字符串
    json_data = json.dumps(data, ensure_ascii=False)
    
    # 加密
    encryption_manager = EncryptionManager()
    encrypted_data = encryption_manager.aes_encrypt(json_data)
    
    return {
        "encrypted_data": encrypted_data
    }

def main():
    print("🔐 AES加密测试数据生成器")
    print("=" * 50)
    
    # 测试数据 - 和你提供的一致
    test_data = {
        "userPhone": "18366542015",
        "userCode": "123456"
    }
    
    print("📋 原始测试数据:")
    print(json.dumps(test_data, ensure_ascii=False, indent=2))
    print()
    
    # 生成加密数据
    print("🔄 正在生成加密数据...")
    encrypted_payload = create_encrypted_payload(test_data.copy())
    
    print("✅ 加密完成!")
    print()
    print("📦 加密载荷:")
    print(json.dumps(encrypted_payload, ensure_ascii=False, indent=2))
    print()
    
    print("🔑 加密数据字符串:")
    print(f'"{encrypted_payload["encrypted_data"]}"')
    print()
    
    # 生成curl命令
    print("🚀 curl测试命令:")
    print("=" * 50)
    
    curl_command = f'''curl -X POST http://124.71.77.44:2233/users/send_sms_code/ \\
  -H "Content-Type: application/json" \\
  -d '{json.dumps(encrypted_payload, ensure_ascii=False)}'
'''
    
    print(curl_command)
    print()
    
    # 保存到文件
    with open('encrypted_curl_command.sh', 'w', encoding='utf-8') as f:
        f.write("#!/bin/bash\n")
        f.write("# AES加密接口测试命令\n")
        f.write("# 生成时间: " + time.strftime('%Y-%m-%d %H:%M:%S') + "\n\n")
        f.write(curl_command)
    
    print("💾 curl命令已保存到: encrypted_curl_command.sh")
    print("🔧 使用方法: chmod +x encrypted_curl_command.sh && ./encrypted_curl_command.sh")
    print()
    
    # 额外信息
    print("ℹ️  测试信息:")
    print(f"   • 时间戳: {int(time.time())}")
    print(f"   • 有效期: 5分钟内")
    print(f"   • 接口地址: http://124.71.77.44:2233/users/send_sms_code/")
    print(f"   • 测试手机号: {test_data['userPhone']}")
    print(f"   • 验证码: {test_data['userCode']}")

if __name__ == "__main__":
    main()
